import pandas as pd

filter_items = ['Average annual expenditures', 
                'Income before taxes',
                'Income after taxes',
                
                ]

use_index = pd.read_excel(r'美国数据\use_index.xlsx', index_col=0)
use_idx = use_index.iloc[:,0].to_list()

dic_us_to_cn = dict(zip(use_index['use_index'],
                    use_index['归并1']))

def save_dict_of_dfs_to_excel(data_dict, output_path):
    with pd.ExcelWriter(output_path) as writer:
        for sheet_name, df in data_dict.items():
            df.to_excel(writer, sheet_name=str(sheet_name))


def read_usa_yearly_data(year):
    data = pd.read_excel(
        f"美国数据\cu-income-deciles-before-taxes-{year}.xlsx", 
        index_col=0,
        header=2
        )
    data.columns = data.columns.str.strip()# 去除列名中的空格
    data.columns = [s.replace('\n','_') for s in data.columns ]
    # data = data.dropna()

    dic = {}
    for name in use_idx:
        # name = use_idx[0]
        if name in data.index.to_list():
            start_n = data.index.to_list().index(name)

            new_idx = data.index.to_list()[start_n:]
            mean_n = new_idx.index('Mean') + start_n
            se = data.iloc[mean_n]
            se.name = name
            dic[name] = se

    df = pd.DataFrame(dic)
    return df.T, data.index.unique().to_list()

def merge_df(df):
    dic = {}
    for idx, se in df.iterrows():
        if idx in dic_us_to_cn:
            new_idx = dic_us_to_cn[idx]
            if new_idx in dic:
                dic[new_idx] += se
            else:
                dic[new_idx] = se
        else:
            dic[idx] = se
    return pd.DataFrame(dic)


def repair_consume_df(df):
    for col, se in df.items():
        numeric_series = pd.to_numeric(se, errors='coerce')
        if numeric_series.isna().any():
            if numeric_series.isna().sum() == 1:
                idx = numeric_series.index[numeric_series.isna()]
                not_na = numeric_series[~numeric_series.isna()]
                filter_idx = [s for s in not_na.index if s not in filter_items]
                not_na_10 = not_na[filter_idx]
                annual_exp = not_na['Average annual expenditures']
                v = annual_exp - not_na_10.sum()
                df.loc[idx, col] = v
                print(col, idx, v)
                df[col] = pd.to_numeric(df[col], errors='coerce')
            else:
                idx = numeric_series.index[numeric_series.isna()]
                not_na = numeric_series[~numeric_series.isna()]
                for idxx in idx:
                    idx_se = df.loc[idxx]
                    idx_se_numeric = pd.to_numeric(idx_se, errors='coerce')
                    if idx_se_numeric.isna().sum() == 1:
                        v_9 = idx_se_numeric[[s for s in idx_se_numeric.index if '10' in s]].sum()
                        v_10 = idx_se_numeric['All_consumer_units']
                        v = 10 * v_10 - v_9
                        df.loc[idxx, col] = v
                print(f'Column {col} contains more than one NaN value')    
    return df



yearly_dic = {}
yearly_df = []
items_yearly = {}
for year in range(2014, 2024):
    df, items = read_usa_yearly_data(year)
    item_dic = {item:1 for item in items}
    items_yearly[year] = item_dic
    df['year'] = year
    yearly_df.append(df)
    yearly_dic[year] = df

df_items = pd.DataFrame(items_yearly)
df_items.to_excel(r'美国数据output\usa_items.xlsx')

df_all_year = pd.concat(yearly_df, axis=0)
save_dict_of_dfs_to_excel(yearly_dic, r'美国数据output\usa_before_repair.xlsx')

repair_df = {}
for k, df in yearly_dic.items():
    repair_df[k] = repair_consume_df(df)
    
save_dict_of_dfs_to_excel(repair_df, r'美国数据output\usa_after_repair1.xlsx')

re_repair_df = {}
cn_repair_df = {}
for k, df in repair_df.items():
    temp_df = repair_consume_df(df)
    re_repair_df[k] = temp_df
    cn_repair_df[k] = merge_df(temp_df)
    
save_dict_of_dfs_to_excel(re_repair_df, r'美国数据output\usa_after_repair_us.xlsx')
save_dict_of_dfs_to_excel(cn_repair_df, r'美国数据output\usa_after_repair_cn.xlsx')